from __future__ import annotations

from typing import Optional
from typing import Tuple

from ppsci.equation import PDE


class Wave(PDE):
    r"""Class for Wave equation.

    $$
    u_{tt} - c^2 \nabla^2 u = 0
    $$

    Args:
        dim (int): Dimension of equation.
        time (bool): Whether to include time dimension.
        c (float): Wave speed.
        detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing.
            Defaults to None.

    Examples:
        >>> import ppsci
        >>> pde = ppsci.equation.Wave(2, True, 1.0)
    """

    def __init__(
        self,
        dim: int = 1,
        time: bool = True,
        c=1.0,
        detach_keys: Optional[Tuple[str, ...]] = None,
    ):
        super().__init__()
        self.detach_keys = detach_keys
        self.dim = dim
        self.time = time

        t, x, y, z = self.create_symbols("t x y z")

        if dim == 1:
            invars = (x,)
        elif dim == 2:
            invars = (x, y)
        elif dim == 3:
            invars = (x, y, z)
        else:
            raise ValueError("Dimension should be 1, 2 or 3.")

        if time:
            invars = (t,) + invars

        u = self.create_function("u", invars)

        if time:
            u_tt = u.diff(t, 2)
            coord_var = invars[1:]
        else:
            u_tt = 0
            coord_var = invars

        poisson = 0
        for invar in coord_var:
            poisson += u.diff(invar, 2)

        self.add_equation("wave", u_tt - c ** 2 * poisson)

        self._apply_detach()
